
#include <iostream>
#include <fstream>
#include <sstream>
#include <unistd.h>
#include <sys/wait.h>

#include "google/protobuf/compiler/code_generator.h"
#include "google/protobuf/descriptor.h"
#include "google/protobuf/compiler/cpp/cpp_generator.h"
#include "babelsnake_generator.h"

#include "config.h"


BabelSnakeGenerator::BabelSnakeGenerator()
{
	/* Declarations */

}

BabelSnakeGenerator::~BabelSnakeGenerator()
{

}


bool BabelSnakeGenerator::handle_field(const ::google::protobuf::FieldDescriptor *descriptor,
								 ::std::string &indent,
								 ::std::stringstream &swigInput,
								  ::std::stringstream &extensionInput) const
{
	/* Declarations */
    ::std::string fieldName = descriptor->name();
    const ::google::protobuf::Descriptor *parentDescriptor;
    ::std::string enumtype;
    char *type;


	//swigInput << indent << descriptor->name() << ";" << ::std::endl;

	swigInput << ::std::endl << indent <<  "// " << fieldName << ::std::endl;

	/* Check if it included has method */
	if(descriptor->is_optional() || descriptor->is_required())
	{
		swigInput << indent << "inline bool has_" << fieldName << "() const;" << ::std::endl;
	}


	swigInput << indent << "inline void clear_" << fieldName << "();" << ::std::endl
	          << indent << "static const int k" << (char)toupper(fieldName[0]) << fieldName.substr(1) << "FieldNumber = " << descriptor->number() << ";" << ::std::endl;



	switch(descriptor->cpp_type())
	{
		/* String type */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_STRING:

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
						  //<< indent << "%rename(" << fieldName << ") mutable_"<< fieldName << ";" << ::std::endl
						  << indent << "inline ::std::string* mutable_" << fieldName << "(int index);" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, const char* value);" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, const char* value, size_t size);" << std::endl
						  << indent << "inline ::std::string* add_" << fieldName << "();" << ::std::endl
						  << indent << "inline void add_" << fieldName << "(const char* value);" << ::std::endl
						  << indent << "inline void add_" << fieldName << "(const char* value, size_t size);" << ::std::endl;


				extensionInput << "%feature(\"shadow\") " << descriptor->containing_type()->name()<< "::mutable_" << fieldName << "(int)" << ::std::endl
							   << "%{" << ::std::endl
							   << "    def " << fieldName << "(self, index):" << ::std::endl
							   << "       if self.sr_size() <= index:" << ::std::endl
							   << "            raise IndexError(\"index out of range\")" << ::std::endl
							   << "       return $action(self, index)" << ::std::endl
							   << "%}" << ::std::endl << ::std::endl;

				extensionInput << "%feature(\"shadow\") " << descriptor->containing_type()->name() << "::set_" << fieldName << "(int, const char*)" << ::std::endl
							   << "%{" << ::std::endl
							   << "    def set_" << fieldName << "(self, index, value):" << ::std::endl
							   << "      if self.sr_size() <= index:" << ::std::endl
							   << "            raise IndexError(\"index out of range\")" << ::std::endl
							   << "      return $action(self, index, value)" << ::std::endl
							   << "%}" << ::std::endl << ::std::endl;

				extensionInput << "%feature(\"shadow\") " << descriptor->containing_type()->name() << "::set_" << fieldName << "(int, const char*, size_t)" << ::std::endl
							   << "%{" << ::std::endl
							   << "    def set_" << fieldName << "(self, index, value, size):" << ::std::endl
							   << "      if self.sr_size() <= index:" << ::std::endl
							   << "            raise IndexError(\"index out of range\")" << ::std::endl
							   << "      return $action(self, index, value, size)" << ::std::endl
							   << "%}" << ::std::endl << ::std::endl;

			/* Singleton */
			}else
			{
				swigInput << indent << "%rename(" << fieldName << ") mutable_"<< fieldName << ";" << ::std::endl
				          << indent << "inline ::std::string* mutable_" << fieldName << "();" << ::std::endl
				    	  << indent << "inline void set_" << fieldName << "(const char* value);" << ::std::endl
					      << indent << "inline void set_" << fieldName << "(const char* value, size_t size);" << std::endl;
			}

			break;

		/* int32 type */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_INT32:

			type = "::google::protobuf::int32";

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
						  << indent << "inline " << type << " " << fieldName << "(int index) const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, " << type << " value);" << ::std::endl
						  << indent << "inline void add_" << fieldName << "(" << type << " value);" << ::std::endl;
			/* Singleton */
			}else
			{
				swigInput << indent << "inline " << type << " " << fieldName << "() const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(" << type << " value);" << ::std::endl;

			}

			break;

		/* int64 type */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_INT64:

			type = "::google::protobuf::int64";

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
						  << indent << "inline " << type << " " << fieldName << "(int index) const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, " << type << " value);" << ::std::endl
				          << indent << "inline void add_" << fieldName << "(" << type << " value);" << ::std::endl;

			/* Singleton */
			}else
			{

				swigInput << indent << "inline " << type << " " << fieldName << "() const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(" << type << ");" << ::std::endl;
			}

			break;

		/* unsigned int32 */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_UINT32:

			type = "::google::protobuf::uint32";

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
						  << indent << "inline " << type << " " << fieldName << "(int index) const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, " << type << " value);" << ::std::endl
						  << indent << "inline void add_" << fieldName << "(" << type << " value);" << ::std::endl;

			/* Singleton */
			}else
			{

				swigInput << indent << "inline " << type << " " << fieldName << "() const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(" << type << ");" << ::std::endl;
			}

			break;

		/* Unsigned int64 */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_UINT64:

			type = "::google::protobuf::uint64";

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
					      << indent << "inline " << type << " " << fieldName << "(int index) const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, " << type << " value);" << ::std::endl
						  << indent << "inline void add_" << fieldName << "(" << type << " value);" << ::std::endl;

			/* Singleton */
			}else
			{

				swigInput << indent << "inline " << type << " " << fieldName << "() const;" << ::std::endl
				     	  << indent << "inline void set_" << fieldName << "(" << type << ");" << ::std::endl;
			}

			break;

		/* Double */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_DOUBLE:

			type = "double";

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
						<< indent << "inline " << type << " " << fieldName << "(int index) const;" << ::std::endl
						<< indent << "inline void set_" << fieldName << "(int index, " << type << " value);" << ::std::endl
						<< indent << "inline void add_" << fieldName << "(" << type << " value);" << ::std::endl;

			/* Singleton */
			}else
			{

				swigInput << indent << "inline " << type << " " << fieldName << "() const;" << ::std::endl
				   	      << indent << "inline void set_" << fieldName << "(" << type << ");" << ::std::endl;
			}

			break;

		/* Float */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_FLOAT:

			type = "float";

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
					   	  << indent << "inline " << type << " " << fieldName << "(int index) const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, " << type << " value);" << ::std::endl
						  << indent << "inline void add_" << fieldName << "(" << type << " value);" << ::std::endl;

			/* Singleton */
			}else
			{

				swigInput << indent << "inline " << type << " " << fieldName << "() const;" << ::std::endl
				          << indent << "inline void set_" << fieldName << "(" << type << ");" << ::std::endl;
			}

			break;

		/* Bool */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_BOOL:

			type = "bool";

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
						  << indent << "inline " << type << " " << fieldName << "(int index) const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, " << type << " value);" << ::std::endl
						  << indent << "inline void add_" << fieldName << "(" << type << " value);" << ::std::endl;

			/* Singleton */
			}else
			{

				swigInput << indent << "inline " << type << " " << fieldName << "() const;" << ::std::endl
				    	  << indent << "inline void set_" << fieldName << "(" << type << ");" << ::std::endl;

			}

			break;

		/* Enum */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_ENUM:

			/* Builds the full type of the message */
			enumtype = descriptor->enum_type()->name();

			for(parentDescriptor = descriptor->containing_type();
				parentDescriptor != NULL;
				parentDescriptor = parentDescriptor->containing_type())
			{
				enumtype.insert(0, parentDescriptor->name() + "_");
			}

			enumtype.insert(0, "::");

			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
						  << indent << "inline " << enumtype << " " << fieldName << "(int index) const;" << ::std::endl
						  << indent << "inline void set_" << fieldName << "(int index, " << enumtype << " value);" << ::std::endl
						  << indent << "inline void add_" << fieldName << "(" << enumtype << " value);" << ::std::endl;

			/* Singleton */
			}else
			{
				swigInput << indent << "inline " << enumtype << " " << fieldName << "() const;" << ::std::endl;
				swigInput << indent << "inline void set_" << fieldName << "(" << enumtype << " value);" << ::std::endl;
			}

			break;

		/* Message */
		case ::google::protobuf::FieldDescriptor::CPPTYPE_MESSAGE:

			/* Builds the full type of the message */
			::std::string type = descriptor->message_type()->name();

			for(parentDescriptor = descriptor->containing_type();
				parentDescriptor != NULL;
				parentDescriptor = parentDescriptor->containing_type())
			{
				type.insert(0, parentDescriptor->name() + "_");
			}

			type.insert(0, "::");


			/* Repeated */
			if(descriptor->is_repeated())
			{
				swigInput << indent << "inline int "<< fieldName << "_size() const;" << ::std::endl
						  << indent << "%rename(" << fieldName << ") mutable_"<< fieldName << ";" << ::std::endl
						  << indent << "inline " << type << "* mutable_" << fieldName << "(int index);" << ::std::endl
						  << indent << "inline " << type << "* add_" << fieldName << "();" << ::std::endl;

			/* Singleton */
			}else
			{
				swigInput << indent << "%rename(" << fieldName << ") mutable_" << fieldName << ";" << ::std::endl
						  << indent << "inline " << type << "* mutable_" << fieldName <<"();" << ::std::endl;
			}

			break;
	}

	/* Return successfully */
	return true;
}

bool BabelSnakeGenerator::handle_enum(const ::google::protobuf::EnumDescriptor *descriptor,
		                        ::std::string &indent,
		                        ::std::stringstream &enumInput) const
{
	/* Declarations */
	::std::string fullName;
	::std::string valueName;
	const ::google::protobuf::Descriptor *parentDescriptor;
	const ::google::protobuf::EnumValueDescriptor *valueDescriptor;

	/* Builds the fullname of the enum */
	fullName = descriptor->name();
	for(parentDescriptor = descriptor->containing_type();
		parentDescriptor != NULL;
		parentDescriptor = parentDescriptor->containing_type())
	{
		fullName.insert(0, parentDescriptor->name()+"_");
	}

	enumInput << indent << "enum " << fullName << ::std::endl
			  << "{" << ::std::endl;


	for(int valueIndex = 0; valueIndex < descriptor->value_count(); valueIndex++)
	{
		valueDescriptor = descriptor->value(valueIndex);

		valueName = fullName + "_" + valueDescriptor->name();

		enumInput << indent << "  " << valueName << " = "
				  << valueDescriptor->number() << "," << ::std::endl;
	}

	enumInput << "};" << ::std::endl;

	return true;
}


bool BabelSnakeGenerator::handle_msg(const ::google::protobuf::Descriptor *descriptor,
							   ::std::string &indent,
							   ::std::stringstream &swigInput,
							   ::std::stringstream &forwardClassInput,
							   ::std::stringstream &enumInput) const
{
	/* Declarations */
	::std::string messageName = descriptor->name();
	const ::google::protobuf::FieldDescriptor *fieldDescriptor;
	const ::google::protobuf::Descriptor *parentDescriptor;
	const ::google::protobuf::Descriptor *nestedDescriptor;
	const ::google::protobuf::EnumDescriptor *enumDescriptor;
	::std::stringstream classInput;
	::std::stringstream extensionInput;
	::std::string nextIndent = indent+"    ";


	/* Handles any enum types */
	for(int enumIndex=0; enumIndex<descriptor->enum_type_count(); enumIndex++)
	{
		handle_enum(descriptor->enum_type(enumIndex), indent, enumInput);
	}

	/* Handles any nested message types */
	for(int nestedIndex=0; nestedIndex<descriptor->nested_type_count(); nestedIndex++)
	{
		handle_msg(descriptor->nested_type(nestedIndex), indent, swigInput,
				   forwardClassInput, enumInput);
	}


	/* Builds the full name of the message class */
	::std::string fullname = messageName;

	for(parentDescriptor = descriptor->containing_type();
		parentDescriptor != NULL;
		parentDescriptor = parentDescriptor->containing_type())
	{
		fullname.insert(0, parentDescriptor->name()+"_");
	}

	forwardClassInput << indent << "class " << fullname << ";" << ::std::endl;

	classInput << indent << "class " << fullname
			  << " : public ::google::protobuf::Message" << ::std::endl
			  << indent << "{" << ::std::endl
			  << indent << "  public:" << ::std::endl

			  /* Constructors and destructors */
			  << indent << "    " << fullname << "();" << ::std::endl
			  << indent << "    " << fullname << "(const " << fullname << "& from);" << ::std::endl
			  << indent << "    virtual ~" << fullname << "();" << ::std::endl
			  << std::endl


			  << nextIndent << "static const " << fullname << "& default_instance();" << ::std::endl

			  << nextIndent << "void Swap(" << fullname << "* other);" << ::std::endl

			  << nextIndent << fullname << "* New() const;" << ::std::endl

			  << nextIndent << "void CopyFrom(const ::google::protobuf::Message& from);" << ::std::endl

			  << nextIndent << "void MergeFrom(const ::google::protobuf::Message& from);" << ::std::endl

			  << nextIndent << "void CopyFrom(const " << fullname << "& from);" << ::std::endl

			  << nextIndent << "void MergeFrom(const " << fullname << "& from);" << ::std::endl

			  << nextIndent << "void Clear();" << ::std::endl

			  << nextIndent << "bool IsInitialized() const;" << ::std::endl

			  << nextIndent << "int ByteSize() const;" << ::std::endl

			  << nextIndent << "int GetCachedSize() const;" << ::std::endl;


	classInput << ::std::endl << nextIndent << "// Nested types" << ::std::endl;
	for(int nestedIndex=0; nestedIndex<descriptor->nested_type_count(); nestedIndex++)
	{
		nestedDescriptor = descriptor->nested_type(nestedIndex);

		/* Builds the full name of the message class */
		::std::string fullname = nestedDescriptor->name();


		for(parentDescriptor = nestedDescriptor->containing_type();
			parentDescriptor != NULL;
			parentDescriptor = parentDescriptor->containing_type())
		{
			fullname.insert(0, parentDescriptor->name()+"_");
		}

		/* Typedef statement */
		classInput << nextIndent << "typedef " << fullname << " "
				  << nestedDescriptor->name() << ";" << ::std::endl;

	}

	for(int enumIndex=0; enumIndex<descriptor->enum_type_count(); enumIndex++)
	{
		enumDescriptor = descriptor->enum_type(enumIndex);

		::std::string fullname = enumDescriptor->name();

		for(parentDescriptor = enumDescriptor->containing_type();
			parentDescriptor != NULL;
			parentDescriptor = parentDescriptor->containing_type())
		{
			fullname.insert(0, parentDescriptor->name()+"_");
		}

		/* Typedef statement */
		classInput << nextIndent << "typedef " << fullname << " "
				  << enumDescriptor->name() << ";" << ::std::endl;
	}



	classInput << ::std::endl << indent << "    " << "// Fields" << ::std::endl;

	/* Iterates over the fields in the message */
	for(int fieldIndex=0; fieldIndex<descriptor->field_count(); fieldIndex++)
	{
		this->handle_field(descriptor->field(fieldIndex), nextIndent, classInput,
						extensionInput);


		//(*this->debugStream) << fieldDescriptor->name() << ::std::endl;
	}

	classInput << indent << "};" << ::std::endl;

	/* Add a __str__ method */
    classInput << indent << "%extend " << fullname << ::std::endl
			  << indent << "{" << ::std::endl
              << indent << "   ::std::string __str__()" << ::std::endl
              << indent << "   {" << ::std::endl
              << indent << "       ::std::stringstream ss(::std::stringstream::in | ::std::stringstream::out);" << ::std::endl
              << indent << "       ss << \"<\" << $self->GetTypeName();" << std::endl
              << indent << "       ss << \" { \" << $self->ShortDebugString() << \" } at \" << $self << \">\";" << ::std::endl
              << indent << "       return ss.str();" << ::std::endl
              << indent << "   }" << ::std::endl
              << indent << "};" << ::std::endl << ::std::endl;


    swigInput << extensionInput.str() << classInput.str();

	/* Return successfully */
	return true;
}

bool BabelSnakeGenerator::Generate(const ::google::protobuf::FileDescriptor *file,
							 const ::std::string &parameter,
							 google::protobuf::compiler::OutputDirectory *output_directory,
							 ::std::string *error) const
{
	/* Declarations */
	::google::protobuf::compiler::cpp::CppGenerator cppGenerator;
	const ::google::protobuf::Descriptor *descriptor;
	const ::google::protobuf::FieldDescriptor *fieldDescriptor;
	::std::stringstream swigInput;
	::std::stringstream forwardClassInput;
	::std::stringstream enumInput;
	::std::string indent("");
	bool result;


	/* Runs the cpp generator */
	result = cppGenerator.Generate(file, parameter, output_directory, error);

	/* Return error from cpp generator */
	if(!result)
	{
		return result;
	}

	::std::cerr << "package: " << file->package() << ::std::endl
		        << "name:    " << file->name() << ::std::endl;

	::std::string module   = file->name().substr(0, file->name().rfind(".proto")) + "_pb";
	::std::string header   = file->name().substr(0, file->name().rfind(".proto")) + ".pb.h";
	::std::string filename = module + ".i";


	/* Handles any enum types */
	for(int enumIndex=0; enumIndex<file->enum_type_count(); enumIndex++)
	{
		handle_enum(file->enum_type(enumIndex), indent, enumInput);
	}

	/* Handles the messages */
	for(int i=0; i<file->message_type_count(); i++)
	{
		 /* Handles the message */
		 handle_msg(file->message_type(i), indent, swigInput, forwardClassInput,
				    enumInput);
	}



	::std::ofstream swigFileStream(filename.c_str(), ::std::ios::out);



	/* Writes the header data */
	swigFileStream  << "%module " << module << ::std::endl << ::std::endl
				    << "%inline" << ::std::endl
				    << "%{"  << ::std::endl << ::std::endl
				    << "#include <sstream>" << ::std::endl
				    << "#include \"google/protobuf/stubs/common.h\"" << ::std::endl
				    << "#include \"google/protobuf/io/coded_stream.h\"" << ::std::endl
				    << "#include \"google/protobuf/message.h\"" << ::std::endl
				    << "#include \"" << header << "\"" << ::std::endl << ::std::endl
				    << "%}" << ::std::endl
				    << ::std::endl
				    << "%include \"basic.i\""
				    << ::std::endl << ::std::endl << ::std::endl;

	/* Writes the forward class declarations */
	swigFileStream << "/* Forward class declarations */" << ::std::endl
			       << forwardClassInput.str() << ::std::endl << ::std::endl;

	/* Writes the enums */
	swigFileStream << "/* Enums */" << ::std::endl
			       << enumInput.str() << ::std::endl << ::std::endl;

    /* Writes the classes */
	swigFileStream << "/* Classes */" << ::std::endl
			       << swigInput.str() << ::std::endl << ::std::endl;

	swigFileStream.close();

	if(!run_swig(filename, error))
	{
		return false;
	}

	/* Returns success */
	return true;
}


bool BabelSnakeGenerator::run_swig(::std::string inputfilename, ::std::string *error) const
{
	/* Declarations */
	pid_t pid;
	int pfd[2];
	int result;


	/* Creates a pipe */
	if(pipe(pfd) != 0)
	{
		error->append("Unable to create a pipe");
		return false;
	}

	/* Create a process */
	if((pid = fork()) < 0)
	{
		error->append("Unable to create process");
		return false;
	}



	/* Child process */
	if(pid == 0)
	{
		close(pfd[0]);
		dup2(pfd[1], STDOUT_FILENO);
		dup2(pfd[1], STDERR_FILENO);

		/* Runs the swig process */
		execl(SWIG_BIN,
				SWIG_BIN,
				PROTOBUF_INC,
				SWIG_INC,
				"-c++",
				"-python",

/* For python 3000 */
#if PYVERSION == 3
				"-py3",
#endif

				inputfilename.c_str(),
				NULL);

		::std::cerr << "Error running swig process";

	}else
	{
		close(pfd[1]);

		/* Waits for swig process to finish */
		waitpid(pid, &result, 0);

		/* Error occurred */
		if(result != 0)
		{
			/* Creates an allocation buffer */
			char *buffer = new char[4096];

			/* Transfers the error message from the swig process to the
			 * error buffer */
			for(int length = read(pfd[0], buffer, 4096);
				length > 0;
				length = read(pfd[0], buffer, 4096))
			{
				error->append(buffer, length);
			}

			/* Frees the buffer */
			delete[] buffer;

			close(pfd[0]);

			/* Return failure */
			return false;

		}else
		{
			close(pfd[0]);
		}
	}

	/* Return success */
	return true;
}


